#!/bin/bash

# YOLO推理服务Docker构建脚本

set -e  # 遇到错误时退出

# 解析命令行参数
DOWNLOAD_MODELS="true"
MODEL_MIRROR="github"
IMAGE_TAG="latest"

while [[ $# -gt 0 ]]; do
    case $1 in
        --no-models)
            DOWNLOAD_MODELS="false"
            shift
            ;;
        --mirror)
            MODEL_MIRROR="$2"
            shift 2
            ;;
        --tag)
            IMAGE_TAG="$2"
            shift 2
            ;;
        -h|--help)
            echo "用法: $0 [选项]"
            echo "选项:"
            echo "  --no-models     跳过模型下载"
            echo "  --mirror MIRROR 使用指定镜像源 (github|gitee)"
            echo "  --tag TAG       指定镜像标签 (默认: latest)"
            echo "  -h, --help      显示帮助信息"
            exit 0
            ;;
        *)
            echo "未知选项: $1"
            echo "使用 --help 查看帮助信息"
            exit 1
            ;;
    esac
done

echo "🐳 开始构建YOLO推理服务Docker镜像..."
echo "📋 构建配置:"
echo "   下载模型: $DOWNLOAD_MODELS"
echo "   镜像源: $MODEL_MIRROR"
echo "   镜像标签: $IMAGE_TAG"

# 检查必要文件
echo "📋 检查必要文件..."
required_files=("Dockerfile" "requirements.txt" "yolo_inference_server.py" "config.yaml" "pip.conf" "download_models.sh")

for file in "${required_files[@]}"; do
    if [ ! -f "$file" ]; then
        echo "❌ 缺少必要文件: $file"
        exit 1
    fi
    echo "✅ 找到文件: $file"
done

# 设置镜像名称
IMAGE_NAME="yolo-inference-service"
FULL_IMAGE_NAME="${IMAGE_NAME}:${IMAGE_TAG}"

echo "🏷️ 镜像名称: $FULL_IMAGE_NAME"

# 构建Docker镜像
echo "🔨 开始构建Docker镜像..."
echo "📦 使用清华源加速pip安装..."
if [ "$DOWNLOAD_MODELS" = "true" ]; then
    echo "🤖 将自动下载YOLO11n和YOLO11s模型..."
else
    echo "⏭️ 跳过模型下载，运行时按需下载..."
fi
echo "⏰ 这可能需要几分钟时间..."

# 构建命令
BUILD_ARGS="--build-arg DOWNLOAD_MODELS=$DOWNLOAD_MODELS --build-arg MODEL_MIRROR=$MODEL_MIRROR"

if docker build $BUILD_ARGS -t "$FULL_IMAGE_NAME" .; then
    echo "✅ Docker镜像构建成功!"
    
    # 显示镜像信息
    echo "📊 镜像信息:"
    docker images "$IMAGE_NAME"
    
    # 显示镜像大小
    IMAGE_SIZE=$(docker images --format "table {{.Size}}" "$FULL_IMAGE_NAME" | tail -n 1)
    echo "📦 镜像大小: $IMAGE_SIZE"
    
    echo ""
    echo "🚀 运行命令:"
    echo "docker run -d -p 8000:8000 --name yolo-service $FULL_IMAGE_NAME"
    echo ""
    echo "🔍 查看日志:"
    echo "docker logs -f yolo-service"
    echo ""
    echo "🛑 停止服务:"
    echo "docker stop yolo-service && docker rm yolo-service"
    
else
    echo "❌ Docker镜像构建失败!"
    echo ""
    echo "🔧 常见问题解决方案:"
    echo "1. 检查网络连接"
    echo "2. 清理Docker缓存: docker system prune -f"
    echo "3. 使用修复版Dockerfile: cp Dockerfile.fixed Dockerfile"
    echo "4. 查看详细错误信息并参考DOCKER_BUILD_TROUBLESHOOTING.md"
    exit 1
fi

echo "🎉 构建完成!"
